#!/usr/bin/env python3
# coding=utf-8

import numpy as np
import cv2

# Roberts 算子
ROBERTS_X = np.array([[-1, 0], [0, 1]], dtype=int)
ROBERTS_Y = np.array([[0, -1], [1, 0]], dtype=int)
PREWITT_X = np.array([[0, 0, -1], [1, 0, -1], [1, 0, -1]])
PREWITT_Y = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])
SOBEL_X = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
SOBEL_Y = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
SCHARR_X = np.array([[-3, 0, 3], [-10, 0, 10], [-3, 0, 3]])  # X方向
SCHARR_Y = np.array([[-3, -10, -3], [0, 0, 0], [3, 10, 3]])  # Y方向

OPERATORS = {
    "roberts": {
        "dx": ROBERTS_X,
        "dy": ROBERTS_Y
    },
    "prewitt": {
        "dx": PREWITT_X,
        "dy": PREWITT_Y
    },
    "sobel": {
        "dx": SOBEL_X,
        "dy": SOBEL_Y
    },
    "scharr": {
        "dx": SCHARR_X,
        "dy": SCHARR_Y
    }
}


def edge_robert(img):
    r, c = img.shape
    new_image = np.zeros((r, c))
    new_imageX = np.zeros((r, c))
    new_imageY = np.zeros((r, c))
    dx = OPERATORS['roberts']['dx']
    dy = OPERATORS['roberts']['dy']
    for i in range(r):
        for j in range(c):
            if (j + 2 <= c) and (i + 2 <= r):
                new_imageX[i, j] = (np.sum(img[i:i + 2, j:j + 2] * dx))**2
                new_imageX[i, j] = (np.sum(img[i:i + 2, j:j + 2] * dy))**2
                new_image[i, j] = (
                    new_imageX[i, j] * new_imageX[i, j] + new_imageY[i, j] * new_imageY[i, j])**0.5
    return np.uint8(new_image)


def edge_prewitt(img):
    r, c = img.shape
    new_image = np.zeros((r, c))
    new_imageX = np.zeros(img.shape)
    new_imageY = np.zeros(img.shape)
    dx = OPERATORS['prewitt']['dx']
    dy = OPERATORS['prewitt']['dy']

    for i in range(r):
        for j in range(c):
            if (j + 3 <= c) and (i + 3 <= r):
                new_imageX[i, j] = (np.sum(img[i:i + 3, j:j + 3] * dx))**2
                new_imageX[i, j] = (np.sum(img[i:i + 3, j:j + 3] * dy))**2
                new_image[i, j] = (
                    new_imageX[i, j] * new_imageX[i, j] + new_imageY[i, j] * new_imageY[i, j])**0.5
    return np.uint8(new_image)


def edge_sobel(img):
    r, c = img.shape
    new_image = np.zeros((r, c))
    new_imageX = np.zeros(img.shape)
    new_imageY = np.zeros(img.shape)
    dx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])  # X方向
    dy = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])  # Y方向

    for i in range(r - 3):
        for j in range(c - 3):
            new_imageX[i + 1, j + 1] = (np.sum(img[i:i + 3, j:j + 3] * dx))**2
            new_imageY[i + 1, j + 1] = (np.sum(img[i:i + 3, j:j + 3] * dy))**2
            new_image[i + 1, j + 1] = (new_imageX[i + 1, j + 1] * new_imageX[i + 1,
                                       j + 1] + new_imageY[i + 1, j + 1] * new_imageY[i + 1, j + 1])**0.5
    return np.uint8(new_image)  # 无方向算子处理的图像


def edge_scharr(img):
    r, c = img.shape
    new_image = np.zeros((r, c))
    new_imageX = np.zeros(img.shape)
    new_imageY = np.zeros(img.shape)
    dx = SCHARR_X
    dy = SCHARR_Y

    for i in range(r - 3):
        for j in range(c - 3):
            new_imageX[i + 1, j + 1] = (np.sum(img[i:i + 3, j:j + 3] * dx))**2
            new_imageY[i + 1, j + 1] = (np.sum(img[i:i + 3, j:j + 3] * dy))**2
            new_image[i + 1, j + 1] = (new_imageX[i + 1, j + 1] * new_imageX[i + 1,
                                       j + 1] + new_imageY[i + 1, j + 1] * new_imageY[i + 1, j + 1])**0.5
    return np.uint8(new_image)  # 无方向算子处理的图像

# 常用的Laplace算子模板 [[0,1,0],[1,-4,1],[0,1,0]] [[1,1,1],[1,-8,1],[1,1,1]]


def edge_laplace(img):
    r, c = img.shape
    new_image = np.zeros((r, c))
    L_sunnzi = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
    # L_sunnzi = np.array([[1,1,1],[1,-8,1],[1,1,1]])

    for i in range(r - 3):
        for j in range(c - 3):
            new_image[i + 1, j +
                      1] = abs(np.sum(img[i:i + 3, j:j + 3] * L_sunnzi))
    return np.uint8(new_image)


def edge_kirsch(img):
    #kirsch算子
    #自定义卷积核，八个方向
    m1 = np.array([[5, 5, 5], [-3, 0, -3], [-3, -3, -3]])
    m2 = np.array([[-3, 5, 5], [-3, 0, 5], [-3, -3, -3]])
    m3 = np.array([[-3, -3, 5], [-3, 0, 5], [-3, -3, 5]])
    m4 = np.array([[-3, -3, -3], [-3, 0, 5], [-3, 5, 5]])
    m5 = np.array([[-3, -3, -3], [-3, 0, -3], [5, 5, 5]])
    m6 = np.array([[-3, -3, -3], [5, 0, -3], [5, 5, -3]])
    m7 = np.array([[5, -3, -3], [5, 0, -3], [5, -3, -3]])
    m8 = np.array([[5, 5, -3], [5, 0, -3], [-3, -3, -3]])

    filterlist = [m1, m2, m3, m4, m5, m6, m7, m8]  # 将各个方向的卷积核放到一起便于统一操作
    # 建立三维数组，第0维表示各个方向卷积后的值
    filtered_list = np.zeros((8, img.shape[0], img.shape[1]))

    for k in range(8):
        # 自定义卷积，其实里面的步骤跟Sobel算子是差不多的
        out = cv2.filter2D(img, cv2.CV_16S, filterlist[k])
        filtered_list[k] = out

    # 取八个方向中的最大值，也就是取第0维的最大值作为图像该点，滤波之后的新的像素值
    final = np.max(filtered_list, axis=0)
    final[np.where(final >= 255)] = 255  # 令像素值大于255的点等于255
    final[np.where(final < 255)] = 0  # 令像素值小于255的点等于0

    return final


edge_detector = {
    "roberts": edge_robert,
    "prewitt": edge_prewitt,
    "sobel": edge_sobel,
    "kirsch": edge_kirsch,
    "scharr": edge_scharr,
    "laplace": edge_laplace,
}


def cv_filter(kernel_x, kernel_y, grayImage):
    # 计算梯度
    x = cv2.filter2D(grayImage, cv2.CV_16S, kernel_x)
    y = cv2.filter2D(grayImage, cv2.CV_16S, kernel_y)

    # 调用convertScaleAbs()函数计算绝对值，
    # 并将图像转换为8位图进行显示，然后进行图像融合
    absX = cv2.convertScaleAbs(x)
    absY = cv2.convertScaleAbs(y)
    res = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
    return res


def cv_scharr(img):
    # Scharr 算子
    x = cv2.Scharr(img, cv2.CV_16S, 1, 0)  # X 方向
    y = cv2.Scharr(img, cv2.CV_16S, 0, 1)  # Y 方向
    absX = cv2.convertScaleAbs(x)
    absY = cv2.convertScaleAbs(y)
    cv_Scharr = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
    return cv_Scharr


def cv_laplace(img):
    # Laplacian
    dst = cv2.Laplacian(img, cv2.CV_16S, ksize=3)
    cv_Laplace = cv2.convertScaleAbs(dst)
    return cv_Laplace


def process(name="roberts", ifile='data/Lena.bmp'):
    edge_fun = edge_detector[name]
    # 图像转化成灰度图像
    img = cv2.imread(ifile, cv2.IMREAD_GRAYSCALE)
    edge = edge_fun(img)
    cv2.imwrite('tmp/edge_lena_' + name + '.png', edge)

    if 'scharr' == name:
        cv_edge = cv_scharr(img)
    elif 'laplace' == name:
        cv_edge = cv_laplace(img)
    elif 'kirsch' == name:
        return
    else:
        # 使用 Numpy 构建卷积核，并对灰度图像在 x 和 y 的方向上做一次卷积运算
        dx = OPERATORS[name]['dx']
        dy = OPERATORS[name]['dy']
        cv_edge = cv_filter(dx, dy, img)
    cv2.imwrite('tmp/edge_lena_' + name + '_cv.png', cv_edge)


if __name__ == "__main__":
    process('roberts')
    process('prewitt')
    process('sobel')
    process('kirsch')
    process('scharr')
    process('laplace')
